/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.transforms.reflect;



import net.bytebuddy.ByteBuddy;

import net.bytebuddy.description.field.FieldDescription;

import net.bytebuddy.description.method.MethodDescription;

import net.bytebuddy.description.modifier.Visibility;

import net.bytebuddy.description.type.TypeDescription;

import net.bytebuddy.description.type.TypeDescription.ForLoadedType;

import net.bytebuddy.description.type.TypeList;

import net.bytebuddy.dynamic.DynamicType;

import net.bytebuddy.dynamic.loading.ClassLoadingStrategy;

import net.bytebuddy.dynamic.scaffold.InstrumentedType;

import net.bytebuddy.dynamic.scaffold.subclass.ConstructorStrategy;

import net.bytebuddy.implementation.ExceptionMethod;

import net.bytebuddy.implementation.FixedValue;

import net.bytebuddy.implementation.Implementation;

import net.bytebuddy.implementation.Implementation.Context;

import net.bytebuddy.implementation.MethodDelegation;

import net.bytebuddy.implementation.bytecode.ByteCodeAppender;

import net.bytebuddy.implementation.bytecode.StackManipulation;

import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;

import net.bytebuddy.implementation.bytecode.Throw;

import net.bytebuddy.implementation.bytecode.assign.Assigner;

import net.bytebuddy.implementation.bytecode.assign.Assigner.Typing;

import net.bytebuddy.implementation.bytecode.assign.TypeCasting;

import net.bytebuddy.implementation.bytecode.constant.IntegerConstant;

import net.bytebuddy.implementation.bytecode.constant.TextConstant;

import net.bytebuddy.implementation.bytecode.member.FieldAccess;

import net.bytebuddy.implementation.bytecode.member.MethodInvocation;

import net.bytebuddy.implementation.bytecode.member.MethodReturn;

import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess;

import net.bytebuddy.jar.asm.Label;

import net.bytebuddy.jar.asm.MethodVisitor;

import net.bytebuddy.jar.asm.Opcodes;

import net.bytebuddy.jar.asm.Type;

import net.bytebuddy.matcher.ElementMatchers;

import com.bff.gaia.unified.sdk.coders.CannotProvideCoderException;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderRegistry;

import com.bff.gaia.unified.sdk.state.Timer;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.splittabledofn.HasDefaultTracker;

import com.bff.gaia.unified.sdk.transforms.splittabledofn.RestrictionTracker;

import com.bff.gaia.unified.sdk.util.UserCodeException;

import com.bff.gaia.unified.sdk.values.TypeDescriptor;

import com.bff.gaia.unified.vendor.guava.com.google.common.primitives.Primitives;

import com.bff.gaia.unified.sdk.util.common.ReflectHelpers;



import javax.annotation.Nullable;

import java.lang.reflect.Constructor;

import java.lang.reflect.InvocationTargetException;

import java.lang.reflect.Method;

import java.util.*;



import static com.bff.gaia.unified.sdk.util.common.ReflectHelpers.findClassLoader;

import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.checkArgument;



/** Dynamically generates a {@link DoFnInvoker} instances for invoking a {@link DoFn}. */

public class ByteBuddyDoFnInvokerFactory implements DoFnInvokerFactory {



  public static final String START_BUNDLE_CONTEXT_PARAMETER_METHOD = "startBundleContext";

  public static final String FINISH_BUNDLE_CONTEXT_PARAMETER_METHOD = "finishBundleContext";

  public static final String PROCESS_CONTEXT_PARAMETER_METHOD = "processContext";

  public static final String ELEMENT_PARAMETER_METHOD = "element";

  public static final String SCHEMA_ELEMENT_PARAMETER_METHOD = "schemaElement";

  public static final String TIMESTAMP_PARAMETER_METHOD = "timestamp";

  public static final String OUTPUT_ROW_RECEIVER_METHOD = "outputRowReceiver";

  public static final String TIME_DOMAIN_PARAMETER_METHOD = "timeDomain";

  public static final String OUTPUT_PARAMETER_METHOD = "outputReceiver";

  public static final String TAGGED_OUTPUT_PARAMETER_METHOD = "taggedOutputReceiver";

  public static final String ON_TIMER_CONTEXT_PARAMETER_METHOD = "onTimerContext";

  public static final String WINDOW_PARAMETER_METHOD = "window";

  public static final String PANE_INFO_PARAMETER_METHOD = "paneInfo";

  public static final String PIPELINE_OPTIONS_PARAMETER_METHOD = "pipelineOptions";

  public static final String RESTRICTION_TRACKER_PARAMETER_METHOD = "restrictionTracker";

  public static final String STATE_PARAMETER_METHOD = "state";

  public static final String TIMER_PARAMETER_METHOD = "timer";



  /**

   * Returns a {@link ByteBuddyDoFnInvokerFactory} shared with all other invocations, so that its

   * cache of generated classes is global.

   */

  public static ByteBuddyDoFnInvokerFactory only() {

    return INSTANCE;

  }



  /**

   * Creates a {@link DoFnInvoker} for the given {@link DoFn} by generating bytecode that directly

   * invokes its methods with arguments extracted from the {@link DoFnInvoker.ArgumentProvider}.

   */

  @Override

  public <InputT, OutputT> DoFnInvoker<InputT, OutputT> invokerFor(DoFn<InputT, OutputT> fn) {

    return newByteBuddyInvoker(fn);

  }



  private static final ByteBuddyDoFnInvokerFactory INSTANCE = new ByteBuddyDoFnInvokerFactory();



  private static final String FN_DELEGATE_FIELD_NAME = "delegate";



  /**

   * A cache of constructors of generated {@link DoFnInvoker} classes, keyed by {@link DoFn} class.

   * Needed because generating an invoker class is expensive, and to avoid generating an excessive

   * number of classes consuming PermGen memory.

   */

  private final Map<Class<?>, Constructor<?>> byteBuddyInvokerConstructorCache =

      new LinkedHashMap<>();



  private ByteBuddyDoFnInvokerFactory() {}



  /** @return the {@link DoFnInvoker} for the given {@link DoFn}. */

  @SuppressWarnings({"unchecked", "rawtypes"})

  public <InputT, OutputT> DoFnInvoker<InputT, OutputT> newByteBuddyInvoker(

      DoFn<InputT, OutputT> fn) {

    return newByteBuddyInvoker(DoFnSignatures.getSignature((Class) fn.getClass()), fn);

  }



  /**

   * Internal base class for generated {@link DoFnInvoker} instances.

   *

   * <p>This class should <i>not</i> be extended directly, or by Unified users. It must be public for

   * generated instances to have adequate access, as they are generated "inside" the invoked {@link

   * DoFn} class.

   */

  public abstract static class DoFnInvokerBase<InputT, OutputT, DoFnT extends DoFn<InputT, OutputT>>

      implements DoFnInvoker<InputT, OutputT> {

    protected DoFnT delegate;



    private Map<String, OnTimerInvoker> onTimerInvokers = new HashMap<>();



    public DoFnInvokerBase(DoFnT delegate) {

      this.delegate = delegate;

    }



    /**

     * Associates the given timer ID with the given {@link OnTimerInvoker}.

     *

     * <p>ByteBuddy does not like to generate conditional code, so we use a map + lookup of the

     * timer ID rather than a generated conditional branch to choose which OnTimerInvoker to invoke.

     *

     * <p>This method has package level access as it is intended only for assembly of the {@link

     * DoFnInvokerBase} not by any subclass.

     */

    void addOnTimerInvoker(String timerId, OnTimerInvoker onTimerInvoker) {

      this.onTimerInvokers.put(timerId, onTimerInvoker);

    }



    @Override

    public void invokeOnTimer(

        String timerId, DoFnInvoker.ArgumentProvider<InputT, OutputT> arguments) {

      @Nullable OnTimerInvoker onTimerInvoker = onTimerInvokers.get(timerId);



      if (onTimerInvoker != null) {

        onTimerInvoker.invokeOnTimer(arguments);

      } else {

        throw new IllegalArgumentException(

            String.format(

                "Attempted to invoke timer %s on %s, but that timer is not registered."

                    + " This is the responsibility of the runner, which must only deliver"

                    + " registered timers.",

                timerId, delegate.getClass().getName()));

      }

    }



    @Override

    public DoFn<InputT, OutputT> getFn() {

      return delegate;

    }

  }



  /** @return the {@link DoFnInvoker} for the given {@link DoFn}. */

  public <InputT, OutputT> DoFnInvoker<InputT, OutputT> newByteBuddyInvoker(

      DoFnSignature signature, DoFn<InputT, OutputT> fn) {

    checkArgument(

        signature.fnClass().equals(fn.getClass()),

        "Signature is for class %s, but fn is of class %s",

        signature.fnClass(),

        fn.getClass());



    try {

      @SuppressWarnings("unchecked")

      DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>> invoker =

          (DoFnInvokerBase<InputT, OutputT, DoFn<InputT, OutputT>>)

              getByteBuddyInvokerConstructor(signature).newInstance(fn);



      for (DoFnSignature.OnTimerMethod onTimerMethod : signature.onTimerMethods().values()) {

        invoker.addOnTimerInvoker(

            onTimerMethod.id(), OnTimerInvokers.forTimer(fn, onTimerMethod.id()));

      }



      return invoker;

    } catch (InstantiationException

        | IllegalAccessException

        | IllegalArgumentException

        | InvocationTargetException

        | SecurityException e) {

      throw new RuntimeException("Unable to bind invoker for " + fn.getClass(), e);

    }

  }



  /**

   * Returns a generated constructor for a {@link DoFnInvoker} for the given {@link DoFn} class.

   *

   * <p>These are cached such that at most one {@link DoFnInvoker} class exists for a given {@link

   * DoFn} class.

   */

  private synchronized Constructor<?> getByteBuddyInvokerConstructor(DoFnSignature signature) {

    Class<? extends DoFn<?, ?>> fnClass = signature.fnClass();

    Constructor<?> constructor = byteBuddyInvokerConstructorCache.get(fnClass);

    if (constructor == null) {

      Class<? extends DoFnInvoker<?, ?>> invokerClass = generateInvokerClass(signature);

      try {

        constructor = invokerClass.getConstructor(fnClass);

      } catch (IllegalArgumentException | NoSuchMethodException | SecurityException e) {

        throw new RuntimeException(e);

      }

      byteBuddyInvokerConstructorCache.put(fnClass, constructor);

    }

    return constructor;

  }



  /** Default implementation of {@link DoFn.SplitRestriction}, for delegation by bytebuddy. */

  public static class DefaultSplitRestriction {

    /** Doesn't split the restriction. */

    @SuppressWarnings("unused")

    public static <InputT, RestrictionT> void invokeSplitRestriction(

        InputT element, RestrictionT restriction, DoFn.OutputReceiver<RestrictionT> receiver) {

      receiver.output(restriction);

    }

  }



  /** Default implementation of {@link DoFn.GetRestrictionCoder}, for delegation by bytebuddy. */

  public static class DefaultRestrictionCoder {

    private final TypeDescriptor<?> restrictionType;



    DefaultRestrictionCoder(TypeDescriptor<?> restrictionType) {

      this.restrictionType = restrictionType;

    }



    /** Doesn't split the restriction. */

    @SuppressWarnings({"unused", "unchecked"})

    public <RestrictionT> Coder<RestrictionT> invokeGetRestrictionCoder(CoderRegistry registry)

        throws CannotProvideCoderException {

      return (Coder) registry.getCoder(restrictionType);

    }

  }



  /** Default implementation of {@link DoFn.NewTracker}, for delegation by bytebuddy. */

  public static class DefaultNewTracker {

    /** Uses {@link HasDefaultTracker} to produce the tracker. */

    @SuppressWarnings("unused")

    public static RestrictionTracker invokeNewTracker(Object restriction) {

      return ((HasDefaultTracker) restriction).newTracker();

    }

  }



  /** Generates a {@link DoFnInvoker} class for the given {@link DoFnSignature}. */

  private static Class<? extends DoFnInvoker<?, ?>> generateInvokerClass(DoFnSignature signature) {

    Class<? extends DoFn<?, ?>> fnClass = signature.fnClass();



    final TypeDescription clazzDescription = new TypeDescription.ForLoadedType(fnClass);



    DynamicType.Builder<?> builder =

        new ByteBuddy()

            // Create subclasses inside the target class, to have access to

            // private and package-private bits

            .with(

                StableInvokerNamingStrategy.forDoFnClass(fnClass)

                    .withSuffix(DoFnInvoker.class.getSimpleName()))



            // class <invoker class> extends DoFnInvokerBase {

            .subclass(DoFnInvokerBase.class, ConstructorStrategy.Default.NO_CONSTRUCTORS)



            //   public <invoker class>(<fn class> delegate) { this.delegate = delegate; }

            .defineConstructor(Visibility.PUBLIC)

            .withParameter(fnClass)

            .intercept(new InvokerConstructor())



            //   public invokeProcessElement(ProcessContext, ExtraContextFactory) {

            //     delegate.<@ProcessElement>(... pass just the right args ...);

            //   }

            .method(ElementMatchers.named("invokeProcessElement"))

            .intercept(new ProcessElementDelegation(clazzDescription, signature.processElement()))



            //   public invokeStartBundle(Context c) { delegate.<@StartBundle>(c); }

            //   ... etc ...

            .method(ElementMatchers.named("invokeStartBundle"))

            .intercept(delegateOrNoop(clazzDescription, signature.startBundle()))

            .method(ElementMatchers.named("invokeFinishBundle"))

            .intercept(delegateOrNoop(clazzDescription, signature.finishBundle()))

            .method(ElementMatchers.named("invokeSetup"))

            .intercept(delegateOrNoop(clazzDescription, signature.setup()))

            .method(ElementMatchers.named("invokeTeardown"))

            .intercept(delegateOrNoop(clazzDescription, signature.teardown()))

            .method(ElementMatchers.named("invokeOnWindowExpiration"))

            .intercept(

                delegateMethodWithExtraParametersOrNoop(

                    clazzDescription, signature.onWindowExpiration()))

            .method(ElementMatchers.named("invokeGetInitialRestriction"))

            .intercept(

                delegateWithDowncastOrThrow(clazzDescription, signature.getInitialRestriction()))

            .method(ElementMatchers.named("invokeSplitRestriction"))

            .intercept(splitRestrictionDelegation(clazzDescription, signature))

            .method(ElementMatchers.named("invokeGetRestrictionCoder"))

            .intercept(getRestrictionCoderDelegation(clazzDescription, signature))

            .method(ElementMatchers.named("invokeNewTracker"))

            .intercept(newTrackerDelegation(clazzDescription, signature.newTracker()));



    DynamicType.Unloaded<?> unloaded = builder.make();



    @SuppressWarnings("unchecked")

    Class<? extends DoFnInvoker<?, ?>> res =

        (Class<? extends DoFnInvoker<?, ?>>)

            unloaded

                .load(

                    ReflectHelpers.findClassLoader(fnClass.getClassLoader()),

                    ClassLoadingStrategy.Default.INJECTION)

                .getLoaded();

    return res;

  }



  private static Implementation getRestrictionCoderDelegation(

      TypeDescription doFnType, DoFnSignature signature) {

    if (signature.processElement().isSplittable()) {

      if (signature.getRestrictionCoder() == null) {

        return MethodDelegation.to(

            new DefaultRestrictionCoder(signature.getInitialRestriction().restrictionT()));

      } else {

        return new DowncastingParametersMethodDelegation(

            doFnType, signature.getRestrictionCoder().targetMethod());

      }

    } else {

      return ExceptionMethod.throwing(UnsupportedOperationException.class);

    }

  }



  private static Implementation splitRestrictionDelegation(

      TypeDescription doFnType, DoFnSignature signature) {

    if (signature.splitRestriction() == null) {

      return MethodDelegation.to(DefaultSplitRestriction.class);

    } else {

      return new DowncastingParametersMethodDelegation(

          doFnType, signature.splitRestriction().targetMethod());

    }

  }



  private static Implementation newTrackerDelegation(

      TypeDescription doFnType, @Nullable DoFnSignature.NewTrackerMethod signature) {

    if (signature == null) {

      // We must have already verified that in this case the restriction type

      // is a subtype of HasDefaultTracker.

      return MethodDelegation.to(DefaultNewTracker.class);

    } else {

      return delegateWithDowncastOrThrow(doFnType, signature);

    }

  }



  /** Delegates to the given method if available, or does nothing. */

  private static Implementation delegateOrNoop(

      TypeDescription doFnType, DoFnSignature.DoFnMethod method) {

    return (method == null)

        ? FixedValue.originType()

        : new DoFnMethodDelegation(doFnType, method.targetMethod());

  }



  /** Delegates method with extra parameters to the given method if available, or does nothing. */

  private static Implementation delegateMethodWithExtraParametersOrNoop(

      TypeDescription doFnType, DoFnSignature.MethodWithExtraParameters method) {

    return (method == null)

        ? FixedValue.originType()

        : new DoFnMethodWithExtraParametersDelegation(doFnType, method);

  }



  /** Delegates to the given method if available, or throws UnsupportedOperationException. */

  private static Implementation delegateWithDowncastOrThrow(

      TypeDescription doFnType, DoFnSignature.DoFnMethod method) {

    return (method == null)

        ? ExceptionMethod.throwing(UnsupportedOperationException.class)

        : new DowncastingParametersMethodDelegation(doFnType, method.targetMethod());

  }



  /**

   * Implements a method of {@link DoFnInvoker} (the "instrumented method") by delegating to a

   * "target method" of the wrapped {@link DoFn}.

   */

  static class DoFnMethodDelegation implements Implementation {

    /** The {@link MethodDescription} of the wrapped {@link DoFn}'s method. */

    protected final MethodDescription targetMethod;

    /** Whether the target method returns non-void. */

    private final boolean targetHasReturn;



    /** Starts {@code null}, initialized by {@link #prepare(InstrumentedType)}. */

    @Nullable

	protected FieldDescription delegateField;



    private final TypeDescription doFnType;



    public DoFnMethodDelegation(TypeDescription doFnType, Method targetMethod) {

      this.doFnType = doFnType;

      this.targetMethod = new MethodDescription.ForLoadedMethod(targetMethod);

      this.targetHasReturn =

          !TypeDescription.VOID.equals(this.targetMethod.getReturnType().asErasure());

    }



    @Override

    public InstrumentedType prepare(InstrumentedType instrumentedType) {

      // Remember the field description of the instrumented type.

      delegateField =

          instrumentedType

              .getSuperClass() // always DoFnInvokerBase

              .getDeclaredFields()

              .filter(ElementMatchers.named(FN_DELEGATE_FIELD_NAME))

              .getOnly();

      // Delegating the method call doesn't require any changes to the instrumented type.

      return instrumentedType;

    }



    @Override

    public ByteCodeAppender appender(final Target implementationTarget) {

      return new ByteCodeAppender() {

        /**

         * @param instrumentedMethod The {@link DoFnInvoker} method for which we're generating code.

         */

        @Override

        public Size apply(

            MethodVisitor methodVisitor,

            Context implementationContext,

            MethodDescription instrumentedMethod) {

          // Figure out how many locals we'll need. This corresponds to "this", the parameters

          // of the instrumented method, and an argument to hold the return value if the target

          // method has a return value.

          int numLocals = 1 + instrumentedMethod.getParameters().size() + (targetHasReturn ? 1 : 0);



          Integer returnVarIndex = null;

          if (targetHasReturn) {

            // Local comes after formal parameters, so figure out where that is.

            returnVarIndex = 1; // "this"

            for (Type param : Type.getArgumentTypes(instrumentedMethod.getDescriptor())) {

              returnVarIndex += param.getSize();

            }

          }



          StackManipulation manipulation =

              new StackManipulation.Compound(

                  // Push "this" (DoFnInvoker on top of the stack)

                  MethodVariableAccess.REFERENCE.loadFrom(0),

                  // Access this.delegate (DoFn on top of the stack)

                  FieldAccess.forField(delegateField).read(),

                  // Cast it to the more precise type

                  TypeCasting.to(doFnType),

                  // Run the beforeDelegation manipulations.

                  // The arguments necessary to invoke the target are on top of the stack.

                  beforeDelegation(instrumentedMethod),

                  // Perform the method delegation.

                  // This will consume the arguments on top of the stack

                  // Either the stack is now empty (because the targetMethod returns void) or the

                  // stack contains the return value.

                  new UserCodeMethodInvocation(returnVarIndex, targetMethod, instrumentedMethod),

                  // Run the afterDelegation manipulations.

                  // Either the stack is now empty (because the instrumentedMethod returns void)

                  // or the stack contains the return value.

                  afterDelegation(instrumentedMethod));



          StackManipulation.Size size = manipulation.apply(methodVisitor, implementationContext);

          return new Size(size.getMaximalSize(), numLocals);

        }

      };

    }



    /**

     * Return the code to the prepare the operand stack for the method delegation.

     *

     * <p>Before this method is called, the stack delegate will be the only thing on the stack.

     *

     * <p>After this method is called, the stack contents should contain exactly the arguments

     * necessary to invoke the target method.

     */

    protected StackManipulation beforeDelegation(MethodDescription instrumentedMethod) {

      return MethodVariableAccess.allArgumentsOf(targetMethod);

    }



    /**

     * Return the code to execute after the method delegation.

     *

     * <p>Before this method is called, the stack will either be empty (if the target method returns

     * void) or contain the method return value.

     *

     * <p>After this method is called, the stack should either be empty (if the instrumented method

     * returns void) or contain the value for the instrumented method to return).

     */

    protected StackManipulation afterDelegation(MethodDescription instrumentedMethod) {

      return new StackManipulation.Compound(

          Assigner.DEFAULT.assign(

              targetMethod.getReturnType(), instrumentedMethod.getReturnType(), Typing.STATIC),

          MethodReturn.of(instrumentedMethod.getReturnType()));

    }

  }



  /**

   * Implements {@link DoFnInvoker}'s method with extra parameters by delegating to a "target

   * method" of the wrapped {@link DoFn}.

   */

  static class DoFnMethodWithExtraParametersDelegation extends DoFnMethodDelegation {

    private final DoFnSignature.MethodWithExtraParameters signature;



    public DoFnMethodWithExtraParametersDelegation(

        TypeDescription clazzDescription, DoFnSignature.MethodWithExtraParameters signature) {

      super(clazzDescription, signature.targetMethod());

      this.signature = signature;

    }



    @Override

    protected StackManipulation beforeDelegation(MethodDescription instrumentedMethod) {

      // Parameters of the wrapper invoker method:

      //   DoFn.ArgumentProvider

      // Parameters of the wrapped DoFn method:

      //   a dynamic set of allowed "extra" parameters in any order subject to

      //   validation prior to getting the DoFnSignature

      ArrayList<StackManipulation> parameters = new ArrayList<>();



      // To load the delegate, push `this` and then access the field

      StackManipulation pushDelegate =

          new StackManipulation.Compound(

              MethodVariableAccess.REFERENCE.loadFrom(0),

              FieldAccess.forField(delegateField).read());



      StackManipulation pushExtraContextFactory = MethodVariableAccess.REFERENCE.loadFrom(1);



      // Push the extra arguments in their actual order.

      for (DoFnSignature.Parameter param : signature.extraParameters()) {

        parameters.add(

            new StackManipulation.Compound(

                pushExtraContextFactory,

                ByteBuddyDoFnInvokerFactory.getExtraContextParameter(param, pushDelegate)));

      }

      return new StackManipulation.Compound(parameters);

    }

  }



  /**

   * Passes parameters to the delegated method by downcasting each parameter of non-primitive type

   * to its expected type.

   */

  private static class DowncastingParametersMethodDelegation extends DoFnMethodDelegation {

    DowncastingParametersMethodDelegation(TypeDescription doFnType, Method method) {

      super(doFnType, method);

    }



    @Override

    protected StackManipulation beforeDelegation(MethodDescription instrumentedMethod) {

      List<StackManipulation> pushParameters = new ArrayList<>();

      TypeList.Generic paramTypes = targetMethod.getParameters().asTypeList();

      for (int i = 0; i < paramTypes.size(); i++) {

        TypeDescription.Generic paramT = paramTypes.get(i);

        pushParameters.add(MethodVariableAccess.of(paramT).loadFrom(i + 1));

        if (!paramT.isPrimitive()) {

          pushParameters.add(TypeCasting.to(paramT));

        }

      }

      return new StackManipulation.Compound(pushParameters);

    }

  }



  /**

   * This wrapper exists to convert checked exceptions to unchecked exceptions, since if this fails

   * the library itself is malformed.

   */

  private static MethodDescription getExtraContextFactoryMethodDescription(

      String methodName, Class<?>... parameterTypes) {

    try {

      return new MethodDescription.ForLoadedMethod(

          DoFnInvoker.ArgumentProvider.class.getMethod(methodName, parameterTypes));

    } catch (Exception e) {

      throw new IllegalStateException(

          String.format(

              "Failed to locate required method %s.%s",

              DoFnInvoker.ArgumentProvider.class.getSimpleName(), methodName),

          e);

    }

  }



  /**

   * Calls a zero-parameter getter on the {@link DoFnInvoker.ArgumentProvider}, which must be on top

   * of the stack.

   */

  private static StackManipulation simpleExtraContextParameter(String methodName) {

    return new StackManipulation.Compound(

        MethodInvocation.invoke(getExtraContextFactoryMethodDescription(methodName)));

  }



  static StackManipulation getExtraContextParameter(

      DoFnSignature.Parameter parameter, final StackManipulation pushDelegate) {



    return parameter.match(

        new DoFnSignature.Parameter.Cases<StackManipulation>() {



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.StartBundleContextParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        START_BUNDLE_CONTEXT_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.FinishBundleContextParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        FINISH_BUNDLE_CONTEXT_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.ProcessContextParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        PROCESS_CONTEXT_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.ElementParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(ELEMENT_PARAMETER_METHOD, DoFn.class)),

                TypeCasting.to(new TypeDescription.ForLoadedType(p.elementT().getRawType())));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.SchemaElementParameter p) {

            ForLoadedType elementType = new ForLoadedType(p.elementT().getRawType());

            ForLoadedType castType =

                elementType.isPrimitive()

                    ? new ForLoadedType(Primitives.wrap(p.elementT().getRawType()))

                    : elementType;



            StackManipulation stackManipulation =

                new StackManipulation.Compound(

                    IntegerConstant.forValue(p.index()),

                    MethodInvocation.invoke(

                        getExtraContextFactoryMethodDescription(

                            SCHEMA_ELEMENT_PARAMETER_METHOD, int.class)),

                    TypeCasting.to(castType));

            if (elementType.isPrimitive()) {

              stackManipulation =

                  new Compound(

                      stackManipulation,

                      Assigner.DEFAULT.assign(

                          elementType.asBoxed().asGenericType(),

                          elementType.asUnboxed().asGenericType(),

                          Typing.STATIC));

            }

            return stackManipulation;

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.TimestampParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        TIMESTAMP_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.TimeDomainParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        TIME_DOMAIN_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.OutputReceiverParameter p) {

            String method =

                p.isRowReceiver() ? OUTPUT_ROW_RECEIVER_METHOD : OUTPUT_PARAMETER_METHOD;

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(method, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.TaggedOutputReceiverParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        TAGGED_OUTPUT_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.OnTimerContextParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        ON_TIMER_CONTEXT_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.WindowParameter p) {

            return new StackManipulation.Compound(

                simpleExtraContextParameter(WINDOW_PARAMETER_METHOD),

                TypeCasting.to(new TypeDescription.ForLoadedType(p.windowT().getRawType())));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.PaneInfoParameter p) {

            return new StackManipulation.Compound(

                pushDelegate,

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(

                        PANE_INFO_PARAMETER_METHOD, DoFn.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.RestrictionTrackerParameter p) {

            // DoFnInvoker.ArgumentProvider.restrictionTracker() returns a RestrictionTracker,

            // but the @ProcessElement method expects a concrete subtype of it.

            // Insert a downcast.

            return new StackManipulation.Compound(

                simpleExtraContextParameter(RESTRICTION_TRACKER_PARAMETER_METHOD),

                TypeCasting.to(new TypeDescription.ForLoadedType(p.trackerT().getRawType())));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.StateParameter p) {

            return new StackManipulation.Compound(

                new TextConstant(p.referent().id()),

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(STATE_PARAMETER_METHOD, String.class)),

                TypeCasting.to(

                    new TypeDescription.ForLoadedType(p.referent().stateType().getRawType())));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.TimerParameter p) {

            return new StackManipulation.Compound(

                new TextConstant(p.referent().id()),

                MethodInvocation.invoke(

                    getExtraContextFactoryMethodDescription(TIMER_PARAMETER_METHOD, String.class)),

                TypeCasting.to(new TypeDescription.ForLoadedType(Timer.class)));

          }



          @Override

          public StackManipulation dispatch(DoFnSignature.Parameter.PipelineOptionsParameter p) {

            return simpleExtraContextParameter(PIPELINE_OPTIONS_PARAMETER_METHOD);

          }

        });

  }



  /**

   * Implements the invoker's {@link DoFnInvoker#invokeProcessElement} method by delegating to the

   * {@link DoFn.ProcessElement} method.

   */

  private static final class ProcessElementDelegation

      extends DoFnMethodWithExtraParametersDelegation {

    private static final MethodDescription PROCESS_CONTINUATION_STOP_METHOD;



    static {

      try {

        PROCESS_CONTINUATION_STOP_METHOD =

            new MethodDescription.ForLoadedMethod(DoFn.ProcessContinuation.class.getMethod("stop"));

      } catch (NoSuchMethodException e) {

        throw new RuntimeException("Failed to locate ProcessContinuation.stop()");

      }

    }



    private final DoFnSignature.ProcessElementMethod signature;



    /** Implementation of {@link MethodDelegation} for the {@link DoFn.ProcessElement} method. */

    private ProcessElementDelegation(

        TypeDescription doFnType, DoFnSignature.ProcessElementMethod signature) {

      super(doFnType, signature);

      this.signature = signature;

    }



    @Override

    protected StackManipulation afterDelegation(MethodDescription instrumentedMethod) {

      if (TypeDescription.VOID.equals(targetMethod.getReturnType().asErasure())) {

        return new StackManipulation.Compound(

            MethodInvocation.invoke(PROCESS_CONTINUATION_STOP_METHOD), MethodReturn.REFERENCE);

      } else {

        return MethodReturn.of(targetMethod.getReturnType().asErasure());

      }

    }

  }



  private static class UserCodeMethodInvocation implements StackManipulation {



    @Nullable

	private final Integer returnVarIndex;

    private final MethodDescription targetMethod;

    private final MethodDescription instrumentedMethod;

    private final TypeDescription returnType;



    private final Label wrapStart = new Label();

    private final Label wrapEnd = new Label();

    private final Label tryBlockStart = new Label();

    private final Label tryBlockEnd = new Label();

    private final Label catchBlockStart = new Label();

    private final Label catchBlockEnd = new Label();



    private final MethodDescription createUserCodeException;



    UserCodeMethodInvocation(

        @Nullable Integer returnVarIndex,

        MethodDescription targetMethod,

        MethodDescription instrumentedMethod) {

      this.returnVarIndex = returnVarIndex;

      this.targetMethod = targetMethod;

      this.instrumentedMethod = instrumentedMethod;

      this.returnType = targetMethod.getReturnType().asErasure();



      boolean targetMethodReturnsVoid = TypeDescription.VOID.equals(returnType);

      checkArgument(

          (returnVarIndex == null) == targetMethodReturnsVoid,

          "returnVarIndex should be defined if and only if the target method has a return value");



      try {

        createUserCodeException =

            new MethodDescription.ForLoadedMethod(

                UserCodeException.class.getDeclaredMethod("wrap", Throwable.class));

      } catch (NoSuchMethodException | SecurityException e) {

        throw new RuntimeException("Unable to find UserCodeException.wrap", e);

      }

    }



    @Override

    public boolean isValid() {

      return true;

    }



    private Object describeType(Type type) {

      switch (type.getSort()) {

        case Type.OBJECT:

          return type.getInternalName();

        case Type.INT:

        case Type.BYTE:

        case Type.BOOLEAN:

        case Type.SHORT:

          return Opcodes.INTEGER;

        case Type.LONG:

          return Opcodes.LONG;

        case Type.DOUBLE:

          return Opcodes.DOUBLE;

        case Type.FLOAT:

          return Opcodes.FLOAT;

        default:

          throw new IllegalArgumentException("Unhandled type as method argument: " + type);

      }

    }



    private void visitFrame(

        MethodVisitor mv, boolean localsIncludeReturn, @Nullable String stackTop) {

      boolean hasReturnLocal = (returnVarIndex != null) && localsIncludeReturn;



      Type[] localTypes = Type.getArgumentTypes(instrumentedMethod.getDescriptor());

      Object[] locals = new Object[1 + localTypes.length + (hasReturnLocal ? 1 : 0)];

      locals[0] = instrumentedMethod.getReceiverType().asErasure().getInternalName();

      for (int i = 0; i < localTypes.length; i++) {

        locals[i + 1] = describeType(localTypes[i]);

      }

      if (hasReturnLocal) {

        locals[locals.length - 1] = returnType.getInternalName();

      }



      Object[] stack = stackTop == null ? new Object[] {} : new Object[] {stackTop};



      mv.visitFrame(Opcodes.F_NEW, locals.length, locals, stack.length, stack);

    }



    @Override

    public Size apply(MethodVisitor mv, Context context) {

      Size size = new Size(0, 0);



      mv.visitLabel(wrapStart);



      String throwableName = new TypeDescription.ForLoadedType(Throwable.class).getInternalName();

      mv.visitTryCatchBlock(tryBlockStart, tryBlockEnd, catchBlockStart, throwableName);



      // The try block attempts to perform the expected operations, then jumps to success

      mv.visitLabel(tryBlockStart);

      size = size.aggregate(MethodInvocation.invoke(targetMethod).apply(mv, context));



      if (returnVarIndex != null) {

        mv.visitVarInsn(Opcodes.ASTORE, returnVarIndex);

        size = size.aggregate(new Size(-1, 0)); // Reduces the size of the stack

      }

      mv.visitJumpInsn(Opcodes.GOTO, catchBlockEnd);

      mv.visitLabel(tryBlockEnd);



      // The handler wraps the exception, and then throws.

      mv.visitLabel(catchBlockStart);

      // In catch block, should have same locals and {Throwable} on the stack.

      visitFrame(mv, false, throwableName);



      // Create the user code exception and throw

      size =

          size.aggregate(

              new Compound(MethodInvocation.invoke(createUserCodeException), Throw.INSTANCE)

                  .apply(mv, context));



      mv.visitLabel(catchBlockEnd);



      // After the catch block we should have the return in scope, but nothing on the stack.

      visitFrame(mv, true, null);



      // After catch block, should have same locals and will have the return on the stack.

      if (returnVarIndex != null) {

        mv.visitVarInsn(Opcodes.ALOAD, returnVarIndex);

        size = size.aggregate(new Size(1, 0)); // Increases the size of the stack

      }

      mv.visitLabel(wrapEnd);

      if (returnVarIndex != null) {

        // Drop the return type from the locals

        mv.visitLocalVariable(

            "res",

            returnType.getDescriptor(),

            returnType.getGenericSignature(),

            wrapStart,

            wrapEnd,

            returnVarIndex);

      }



      return size;

    }

  }



  /**

   * A constructor {@link Implementation} for a {@link DoFnInvoker class}. Produces the byte code

   * for a constructor that takes a single argument and assigns it to the delegate field.

   */

  private static final class InvokerConstructor implements Implementation {

    @Override

    public InstrumentedType prepare(InstrumentedType instrumentedType) {

      return instrumentedType;

    }



    @Override

    public ByteCodeAppender appender(final Target implementationTarget) {

      return (methodVisitor, implementationContext, instrumentedMethod) -> {

        StackManipulation.Size size =

            new StackManipulation.Compound(

                    // Load the this reference

                    MethodVariableAccess.REFERENCE.loadFrom(0),

                    // Load the delegate argument

                    MethodVariableAccess.REFERENCE.loadFrom(1),

                    // Invoke the super constructor (default constructor of Object)

                    MethodInvocation.invoke(

                        new TypeDescription.ForLoadedType(DoFnInvokerBase.class)

                            .getDeclaredMethods()

                            .filter(

                                ElementMatchers.isConstructor()

                                    .and(ElementMatchers.takesArguments(DoFn.class)))

                            .getOnly()),

                    // Return void.

                    MethodReturn.VOID)

                .apply(methodVisitor, implementationContext);

        return new ByteCodeAppender.Size(size.getMaximalSize(), instrumentedMethod.getStackSize());

      };

    }

  }

}